#include "znx_tcp_proc.h"
#include "znx_worker_proc.h"

#define MAX_DO_ACCEPT_ONCE      32

char *znx_tcp_accept_thread_tag = "tcp_accepter_thread_tag";

znx_tcp_listener_t g_tcp_listener_list = {
    .next = NULL,
};

znx_tcp_accepter_t g_tcp_accepter = {
    0, NULL, NULL
};

znx_tcp_connect_opt_t default_tcp_conn_opt = {
    .nodelay = 1,
    .quickack = 1,
};


static znx_bool_t
znx_tcp_listener_open(znx_tcp_listener_t *l)
{
    znx_syscall_result_t res;
    res = znx_syscall_socket(l->address.sockaddr.sa_family, SOCK_STREAM, 0);
    if (res.eno != 0) {
        ZNX_DEF_LOG_ERROR("socket() %s failed, errno: %d", l->listen_addr, res.eno);
        return ZNX_FALSE;
    }

    l->sockfd = (int)res.res;

    int opt_val = 1;

    if (l->reuseport) {
        res = znx_syscall_setsockopt(l->sockfd, SOL_SOCKET, SO_REUSEPORT,
            &opt_val, sizeof(int));
        if (res.eno != 0) {
            ZNX_DEF_LOG_ERROR("setsockopt(SO_REUSEPORT) %s failed, errno: %d",
                l->listen_addr, res.eno);
            return ZNX_FALSE;
        }
    }

    if (l->reuseaddr) {
        res = znx_syscall_setsockopt(l->sockfd, SOL_SOCKET, SO_REUSEADDR,
            &opt_val, sizeof(int));
        if (res.eno != 0) {
            ZNX_DEF_LOG_ERROR("setsockopt(SO_REUSEADDR) %s failed, errno: %d",
                l->listen_addr, res.eno);
            return ZNX_FALSE;
        }
    }

    if (l->ipv6only && l->address.sockaddr.sa_family == AF_INET6) {
        res = znx_syscall_setsockopt(l->sockfd, IPPROTO_IPV6, IPV6_V6ONLY,
            &opt_val, sizeof(int));
        if (res.eno != 0) {
            ZNX_DEF_LOG_ERROR("setsockopt(IPV6_V6ONLY) %s failed, errno: %d",
                l->listen_addr, res.eno);
            return ZNX_FALSE;
        }
    }

    res = znx_syscall_set_nonblock(l->sockfd);
    if (res.eno != 0) {
        ZNX_DEF_LOG_ERROR("setnonblock %s failed, errno: %d", l->listen_addr, res.eno);
        return ZNX_FALSE;
    }

#if 0
    res = znx_syscall_set_close_exec(l->sockfd);
    if (res.eno != 0) {
        ZNX_DEF_LOG_ERROR("setcloseexec %s failed, errno: %d", l->listen_addr, res.eno);
        return ZNX_FALSE;
    }
#endif

    res = znx_syscall_bind(l->sockfd, &l->address.sockaddr,
        (socklen_t)znx_address_socklen(&l->address));
    if (res.eno != 0) {
        ZNX_DEF_LOG_ERROR("bind %s failed, errno: %d", l->listen_addr, res.eno);
        return ZNX_FALSE;
    }

    if (l->address.sockaddr.sa_family == AF_UNIX) {
        mode_t mode = (S_IRUSR|S_IWUSR|S_IRGRP|S_IWGRP|S_IROTH|S_IWOTH);
        if (chmod(l->address.sockaddr_un.sun_path, mode) == -1) {
            ZNX_DEF_LOG_ERROR("chmod unix: %s failed, errno: %d",
                l->address.sockaddr_un.sun_path, errno);
            return ZNX_FALSE;
        }
    }

    res = znx_syscall_listen(l->sockfd, l->backlog);
    if (res.eno != 0) {
        ZNX_DEF_LOG_ERROR("listen() to %s failed, errno: %d", l->listen_addr, res.eno);
        return ZNX_FALSE;
    }

    return ZNX_TRUE;
}


// run in main thread.
znx_bool_t
znx_tcp_listeners_open()
{
    znx_tcp_listener_t          *listener;
    listener = g_tcp_listener_list.next;
    while (listener != NULL) {
        if (!znx_tcp_listener_open(listener)) {
            ZNX_DEF_LOG_ERROR("open %s listener failed, exit process", listener->listen_addr);
            return ZNX_FALSE;
        }
        listener = listener->next;
    }

    return ZNX_TRUE;
}


static void
znx_tcp_listener_close(znx_tcp_listener_t *l)
{
    if (l->sockfd == INVALID_FD) {
        ZNX_DEF_LOG_ERROR("close invalid tcp listener");
        return;
    }

    znx_syscall_close(l->sockfd);
}


// run in main thread.
void
znx_tcp_listeners_close()
{
    znx_tcp_listener_t  *listener;
    listener = g_tcp_listener_list.next;
    while (listener != NULL) {
        znx_tcp_listener_close(listener);
        listener = listener->next;
    }
}


// run in worker thread.
static void
znx_tcp_accept_dispatch_handler(znx_dispatch_event_t *dispatch_event)
{
    znx_tcp_listener_t *ls;
    ls = dispatch_event->save_data_1;
    znx_conn_t *conn = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn, ZNX_CONN_EPOLL_ET);
    conn->fd = (int)(intptr_t)dispatch_event->data;
    conn->read.handler = ls->accept_handler;
    conn->ls = ls;

    // see znx_tcp_accept_dispatch_event_alloc.
    free(dispatch_event);

    // notify using read event.
    conn->read.handler(&conn->read);
}


// run in accept thread.
static znx_dispatch_event_t *
znx_tcp_accept_dispatch_event_alloc(int sockfd, znx_tcp_listener_t *ls)
{
    znx_dispatch_event_t    *dispatch_event;
    // the znx_thread_malloc interface is not used here
    // because there is no way to reuse memory.
    // it is inefficient to apply for memory from the global every time.
    dispatch_event = malloc(sizeof(znx_dispatch_event_t));
    dispatch_event->data = (void *)(intptr_t)sockfd;
    dispatch_event->save_data_1 = ls;
    dispatch_event->handler = znx_tcp_accept_dispatch_handler;

    return dispatch_event;
}


// run in accept thread.
static znx_bool_t
znx_tcp_accept_sockfd_set(int accept_fd, znx_tcp_listener_t *ls)
{
    int                             opt_val;
    znx_syscall_result_t             res;

    if (ls->nodelay) {
        res = znx_syscall_setsockopt(accept_fd,
            IPPROTO_TCP, TCP_NODELAY, &opt_val, sizeof(int));
        if (res.eno != 0) {
            ZNX_DEF_LOG_ERROR("setsockopt(TCP_NODELAY) %d failed, errno: %d",
                accept_fd, res.eno);
            return ZNX_FALSE;
        }
    }

    if (ls->quickack) {
        res = znx_syscall_setsockopt(accept_fd,
            IPPROTO_TCP, TCP_QUICKACK, &opt_val, sizeof(int));
        if (res.eno != 0) {
            ZNX_DEF_LOG_ERROR("setsockopt(TCP_QUICKACK) %d failed, errno: %d",
                accept_fd, res.eno);
            return ZNX_FALSE;
        }
    }

    return ZNX_TRUE;
}


// run in accept thread.
static void
znx_tcp_accept_handle(znx_event_t *event)
{
    znx_conn_t                      *c;
    znx_tcp_listener_t              *ls;
    znx_syscall_result_t            res;
    int                             accept_fd;
    znx_dispatch_event_t            *dispatch_event;

    c = event->data;
    ls = c->data;

    for (int i = 0; i < MAX_DO_ACCEPT_ONCE; i++) {
        res = znx_syscall_accept4(c->fd, NULL, NULL, SOCK_NONBLOCK|SOCK_CLOEXEC);
        if (res.eno != 0) {
            break; // EAGAIN or other error
        }

        accept_fd = (int)res.res;
        if (!znx_tcp_accept_sockfd_set(accept_fd, ls)) {
            ZNX_DEF_LOG_ERROR("set accept fd %d tcp options failed", accept_fd);
            znx_syscall_close(accept_fd);
            continue;
        }

        dispatch_event =
            znx_tcp_accept_dispatch_event_alloc(accept_fd, ls);
        znx_worker_processor_post_event(dispatch_event);
    }
}


typedef struct {
    int                         n_listener;
    znx_conn_t                  *accepter_conns;
} znx_tcp_accepter_loop_data_t;


static inline void
znx_tcp_accepter_loop_data_clean(void *data)
{
    znx_conn_t   *conn;
    znx_tcp_accepter_loop_data_t *accepter_loop_data = data;

    if (accepter_loop_data->accepter_conns == NULL) {
        ZNX_DEF_LOG_FATAL("unreachable");
    }

    for (int i = 0; i < accepter_loop_data->n_listener; i++) {
        conn = accepter_loop_data->accepter_conns + i;
        znx_event_loop_remove_conn(conn);
        znx_conn_close(conn);
    }

    znx_thread_free(accepter_loop_data->accepter_conns);
    znx_thread_free(accepter_loop_data);
}


// run in accept thread.
static void
znx_tcp_accepter_loop_init(znx_event_loop_data_t *loop_data)
{
    znx_conn_t              *conn, *conns;
    znx_tcp_listener_t      *listener;
    int                     i = 0;
    int                     n = 0;

    listener = g_tcp_listener_list.next;
    while (listener != NULL) {
        n++;
        listener = listener->next;
    }

    znx_tcp_accepter_loop_data_t *accepter_loop_data;
    accepter_loop_data = znx_thread_calloc(sizeof(znx_tcp_accepter_loop_data_t));
    conns = znx_thread_calloc(sizeof(znx_conn_t) * (size_t)n);

    listener = g_tcp_listener_list.next;
    for (i = 0; i < n; i++) {
        conn = conns + i;
        znx_conn_init(conn, ZNX_CONN_EPOLL_LT);
        conn->fd = listener->sockfd;
        conn->shared = 1;
        conn->data = listener;
        conn->read.handler = znx_tcp_accept_handle;

        if (!znx_event_loop_start_io(conn)) {
            ZNX_DEF_LOG_FATAL("start tcp listener %s io failed", listener->listen_addr);
        }

        if (!znx_event_loop_enable_read(conn)) {
            ZNX_DEF_LOG_FATAL("enabled tcp listener %s io failed", listener->listen_addr);
        }

        listener = listener->next;
    }

    accepter_loop_data->n_listener = n;
    accepter_loop_data->accepter_conns = conns;
    znx_event_loop_cleanup_add(znx_tcp_accepter_loop_data_clean,
        accepter_loop_data);
}


// run in main thread.
void
znx_tcp_accepter_start(int nthreads)
{
    if (nthreads <= 0) {
        ZNX_DEF_LOG_FATAL("tcp accept threads must be greate than 0");
    }

    if (g_tcp_listener_list.next == NULL) {
        ZNX_DEF_LOG_WARN("no tcp listener");
        return;
    }

    znx_event_loop_data_t       *datas;
    znx_thread_base_t           *threads;

    datas = znx_thread_calloc(
        sizeof(znx_event_loop_data_t) * (size_t)nthreads);
    threads = znx_thread_calloc(
        sizeof(znx_thread_base_t) * (size_t)nthreads);

    char thread_name[64] = {0};

    for (int i = 0; i < nthreads; i++) {
        snprintf(thread_name, 64, "tcp-accepter-%d", i);

        znx_event_loop_data_init(
            datas+i,
            znx_tcp_accepter_loop_init,
            NULL,
            NULL);

        znx_thread_base_start(
            znx_tcp_accept_thread_tag,
            threads+i,
            znx_event_loop_runner,
            datas+i,
            0, 0);

        znx_thread_set_name((threads+i)->tid, thread_name);
    }

    g_tcp_accepter.nthreads = nthreads;
    g_tcp_accepter.threads = threads;
    g_tcp_accepter.datas = datas;
}


// run in main thread.
void
znx_tcp_accepter_stop()
{
    if (g_tcp_accepter.nthreads == 0) {
        return;
    }

    int                 i;
    for (i = 0; i < g_tcp_accepter.nthreads; i++) {
        znx_event_loop_quit(g_tcp_accepter.threads+i);
    }

    for (i = 0; i < g_tcp_accepter.nthreads; i++) {
        znx_thread_base_join(g_tcp_accepter.threads+i, NULL);
    }

    znx_thread_free(g_tcp_accepter.datas);
    znx_thread_free(g_tcp_accepter.threads);
}



////////////////////////////// tcp connect //////////////////////////////

typedef struct {
    znx_tcp_connect_handler_pt  handler;
    void                        *data;
} znx_tcp_connect_data_t;


// run in worker threads.
static void
znx_tcp_connect_done(znx_conn_t *conn, znx_result_t res)
{
    znx_tcp_connect_handler_pt  handler;
    void                        *data;

    znx_tcp_connect_data_t *conn_data = conn->data;
    handler = conn_data->handler;
    data = conn_data->data;
    znx_thread_free(conn_data);

    conn->data = NULL;
    conn->write.handler = NULL;

    if (res != ZNX_SUCCESS) {
        if (conn->fd != INVALID_FD) {
            znx_conn_close(conn);
        }
        znx_thread_free(conn);
        conn = NULL;
    }

    handler(data, conn, res);
}


// run in worker threads.
static void
znx_tcp_conn_wait_connect_handler(znx_event_t *event)
{
    znx_conn_t *conn = event->data;
    znx_event_loop_remove_conn(conn);

    if (event->timedout) {
        ZNX_DEF_LOG_DEBUG("connect timedout, fd: %d", conn->fd);
        event->timedout = 0;
        znx_tcp_connect_done(conn, ZNX_TIMEOUT);
        return;
    }

    int opt_val     = 0;
    int opt_size    = sizeof(opt_val);

    znx_syscall_result_t res;
    res = znx_syscall_getsockopt(conn->fd, SOL_SOCKET, SO_ERROR, &opt_val, (socklen_t *)&opt_size);
    if (res.eno != 0 || opt_val != 0) {
        ZNX_DEF_LOG_DEBUG("connect failed, fd: %d, errno: %d, opt_val: %d",
            conn->fd, res.res, opt_val);
        znx_tcp_connect_done(conn, ZNX_FAILED);
        return;
    }

    znx_tcp_connect_done(conn, ZNX_SUCCESS);
}


// run in worker threads.
static void
znx_tcp_connect_setup_helper(znx_conn_t *conn, znx_address_t *address,
    znx_address_t *bind_address, int64_t timeout_ms, znx_tcp_connect_opt_t *opts)
{
    if (!current_thread_base) {
        ZNX_DEF_LOG_FATAL("current thread is not event_loop");
    }

    if (opts == NULL) {
        opts = &default_tcp_conn_opt;
    }

    znx_syscall_result_t res;
    res = znx_syscall_socket(address->sockaddr.sa_family, SOCK_STREAM, 0);
    if (res.eno != 0) {
        ZNX_DEF_LOG_ERROR("socket() failed, errno: %d", res.eno);
        return znx_tcp_connect_done(conn, ZNX_FAILED);
    }
    conn->fd = (int)res.res;

    int opt_val = 1;
    if (opts->nodelay) {
        res = znx_syscall_setsockopt(conn->fd, IPPROTO_TCP, TCP_NODELAY, &opt_val, sizeof(int));
        if (res.eno != 0) {
            ZNX_DEF_LOG_ERROR("setsockopt(TCP_NODELAY) %d failed, errno: %d", conn->fd, res.eno);
            return znx_tcp_connect_done(conn, ZNX_FAILED);
        }
    }

    if (opts->quickack) {
        res = znx_syscall_setsockopt(conn->fd, IPPROTO_TCP, TCP_QUICKACK, &opt_val, sizeof(int));
        if (res.eno != 0) {
            ZNX_DEF_LOG_ERROR("setsockopt(TCP_QUICKACK) %d failed, errno: %d", conn->fd, res.eno);
            return znx_tcp_connect_done(conn, ZNX_FAILED);
        }
    }

    res = znx_syscall_set_nonblock(conn->fd);
    if (res.eno != 0) {
        ZNX_DEF_LOG_ERROR("setnonblock %d failed, errno: %d", conn->fd, res.eno);
        return znx_tcp_connect_done(conn, ZNX_FAILED);
    }

    res = znx_syscall_set_close_exec(conn->fd);
    if (res.eno != 0) {
        ZNX_DEF_LOG_ERROR("setcloseexec %d failed, errno: %d", conn->fd, res.eno);
        return znx_tcp_connect_done(conn, ZNX_FAILED);
    }

    if (bind_address != NULL) {
        res = znx_syscall_bind(conn->fd, &bind_address->sockaddr,
            (socklen_t)znx_address_socklen(bind_address));
        if (res.eno != 0) {
            ZNX_DEF_LOG_ERROR("bind %d local address failed, errno: %d", conn->fd, res.eno);
            return znx_tcp_connect_done(conn, ZNX_FAILED);
        }
    }

    res = znx_syscall_connect(conn->fd, &address->sockaddr, (socklen_t)znx_address_socklen(address));
    // connect success
    if (res.eno == 0) {
        return znx_tcp_connect_done(conn, ZNX_SUCCESS);
    }

    if (res.eno != EINPROGRESS) {
        ZNX_DEF_LOG_ERROR("connect() failed, fd: %d, errno: 5d", conn->fd, res.eno);
        return znx_tcp_connect_done(conn, ZNX_FAILED);
    }

    conn->write.handler = znx_tcp_conn_wait_connect_handler;

    if (timeout_ms > 0) {
        znx_event_loop_add_timer(&conn->write, timeout_ms);
    }

    znx_event_loop_start_io(conn);
    znx_event_loop_enable_write(conn);
    return; // wait event loop callback
}


// run in worker threads.
void znx_tcp_connect_setup(
    znx_address_t *address,
    znx_address_t *bind_address,
    int64_t timeout_ms,
    znx_tcp_connect_opt_t *opts,
    znx_tcp_connect_handler_pt connect_handler,
    void *data)
{
    assert(current_thread_is_worker_thread());

    znx_tcp_connect_data_t *connect_data =
        znx_thread_calloc(sizeof(znx_tcp_connect_data_t));
    connect_data->handler = connect_handler;
    connect_data->data = data;

    znx_conn_t *conn = znx_thread_calloc(sizeof(znx_conn_t));
    znx_conn_init(conn, ZNX_CONN_EPOLL_ET);
    conn->data = connect_data;

    znx_tcp_connect_setup_helper(conn, address, bind_address, timeout_ms, opts);
}